# %%
import math
from datetime import datetime, timedelta
from typing import List
import matplotlib.pyplot as plt  # 导入绘图库
import platform
# 设置中文字体
if platform.system().lower() == "windows":
    plt.rcParams["font.family"] = "SimHei"
else:
    plt.rcParams["font.sans-serif"] = ["Arial Unicode MS"]  
plt.rcParams["axes.unicode_minus"] = False


class Curve:
    def __init__(
        self,
        term=0.0,
        ytm=0.0,
        df=0.0,
        zero_rate=0.0,
        zero_rate_continue=0.0,
        maturity_date=None,
    ):
        self.term = term
        self.ytm = ytm
        self.df = df
        self.zero_rate = zero_rate
        self.zero_rate_continue = zero_rate_continue
        self.maturity_date = maturity_date


def cubic_hermite_interpolation(x, x0, x1, y0, y1, dy0, dy1):
    """
    三次埃米特尔插值
    :param x: 插值点
    :param x0: 左端点
    :param x1: 右端点
    :param y0: 左端点的函数值
    :param y1: 右端点的函数值
    :param dy0: 左端点的导数值
    :param dy1: 右端点的导数值
    :return: 插值结果
    """
    # 计算插值区间长度
    dx = x1 - x0
    # 计算插值点相对于左端点的偏移量
    t = (x - x0) / dx
    # 三次埃米特尔插值公式
    h00 = 2 * t**3 - 3 * t**2 + 1
    h10 = t**3 - 2 * t**2 + t
    h01 = -2 * t**3 + 3 * t**2
    h11 = t**3 - t**2
    # 计算结果
    return h00 * y0 + h10 * dx * dy0 + h01 * y1 + h11 * dx * dy1


def completion_curve(curve_list: List[Curve], add_list: List[float]):
    """插值补全每个 1/frequence 年的到期收益率"""
    for d in add_list:
        for i in range(len(curve_list)):
            curve = curve_list[i]
            if curve.term == d:
                break
            if curve.term < d:
                continue
            if curve.term > d:
                # 获取插值区间内的两个端点
                x0 = curve_list[i - 1].term
                x1 = curve.term
                y0 = curve_list[i - 1].ytm
                y1 = curve.ytm
                # 计算导数（斜率），这里使用简单的差分近似
                dy0 = (curve_list[i].ytm - curve_list[i - 1].ytm) / (
                    curve_list[i].term - curve_list[i - 1].term
                )
                dy1 = (curve_list[i + 1].ytm - curve_list[i].ytm) / (
                    curve_list[i + 1].term - curve_list[i].term
                )
                # 使用三次埃米特尔插值
                rate = cubic_hermite_interpolation(d, x0, x1, y0, y1, dy0, dy1)
                completion_curve = Curve(term=d, ytm=rate)
                curve_list.append(completion_curve)
                break
    curve_list.sort(key=lambda x: x.term)
    return curve_list


def bootstrapping_curve(curve_list: List[Curve], frequence: int, calc_date: datetime):
    """票息剥离构建曲线"""
    zero_curve_list = []
    for curve in curve_list:
        zero_curve = Curve()
        term = curve.term
        rate = curve.ytm
        zero_curve.term = term
        zero_curve.ytm = rate
        # 到期日，乘以12四舍五入，加上月份，这里暂不考虑节假日等因素，以及非月份等时间。
        month = round(term * 12)
        zero_curve.maturity_date = calc_date + timedelta(
            days=month * 30
        )  # 简单假设每个月30天
        zero_rate = rate
        if term <= 1:
            df = 1 / (1 + rate * term)
            zero_curve.df = df
            zero_curve.zero_rate = rate
        else:
            zero_rate = cal_spot_rate(zero_curve_list, term, rate, frequence)
            zero_curve.zero_rate = zero_rate
            zero_curve.df = 1 / math.pow(1 + zero_rate, zero_curve.term * frequence)
        zero_rate_continue = frequence * math.log(1 + zero_rate / frequence)
        zero_curve.zero_rate_continue = zero_rate_continue
        print(
            f"期限：{zero_curve.term:.6f}，到期收益率{zero_curve.ytm:.6f}，折现{zero_curve.df:.6f}，零息{zero_curve.zero_rate:.6f}，零息连续复利{zero_rate_continue:.6f}"
        )
        zero_curve_list.append(zero_curve)
    return zero_curve_list


def list_term(begin: float, end: float, frequence: float, scale: int):
    """从1到50年中，按照frequence，构造出每个点"""
    double_list = []
    term = begin
    double_list.append(begin)
    step = 1.0 / frequence
    while end - term >= 0.001:  # 避免误差
        double_list.append(round(term, scale))
        term += step
    return double_list


def bootstrapping1(
    curve_list: List[Curve], term: float, ytm: float, x: float, frequence: int
):
    """传入最后期限点的ytm和spotRate算出折现后100，该方法用于牛顿逼近"""
    sum_val = 0
    c = 100 * ytm / frequence
    for curve in curve_list:
        if curve.term < 1 and curve.term != 1.0 / frequence:
            # 把0.08333,0.25,0.75的过滤掉
            continue
        if curve.term <= 1:
            # 单利
            sum_val += c / (1 + curve.zero_rate * curve.term)
        elif curve.term < term:
            sum_val += c / math.pow(
                1 + curve.zero_rate / frequence, curve.term * frequence
            )
    sum_val += (100 + c) / math.pow(1 + x / frequence, term * frequence)
    return sum_val


def bootstrapping2(curve_list: List[Curve], term: float, ytm: float, step: float):
    """这个是知乎大神HighFinance的剥息法，更为简单。效率也比牛顿法好"""
    sum_val = 0
    for curve in curve_list:
        if curve.term >= term:
            break
        if curve.term < 1 and curve.term != step:
            continue
        sum_val += curve.df
    c = step * ytm
    return (1 - 1 * c * sum_val) / (1 + c)


def cal_spot_rate(curve_list: List[Curve], term: float, ytm: float, frequence: int):
    """用牛顿法算spot rate"""

    def f(x):
        return 100 - bootstrapping1(curve_list, term, ytm, x, frequence)

    return root_newton(f, 0.01)


def root_newton(f, initial_guess, tol=1e-6, max_iter=1000):
    """牛顿法求根"""
    x = initial_guess
    for _ in range(max_iter):
        fx = f(x)
        if abs(fx) < tol:
            return x
        dfx = (f(x + tol) - fx) / tol  # 数值微分
        x = x - fx / dfx
    return x


def test_curve_list():
    """测试数据"""
    curve_list = []
    # chinamoney 2023-5-10 测试数据，保留关键期限点
    terms = [
        # 0,
        # 0.083,
        # 0.25,
        0.5,
        0.75,
        1.0,
        # 2.0,
        3.0,
        4.0,
        # 5.0,
        6.0,
        7.0,
        8.0,
        # 9.0,
        10.0,
        15.0,
        20.0,
        30.0,
        # 40.0,
        # 50.0,
    ]
    ytm = [
        # 1.4734,
        # 1.5956,
        # 1.6,
        1.6241,
        1.5659,
        1.5547,
        # 1.507,
        1.5923,
        1.6394,
        # 1.6672,
        1.7311,
        1.7773,
        1.7922,
        # 1.7974,
        1.8002,
        1.9424,
        2.0934,
        1.975,
        # 2.0643,
        # 2.12,
    ]
    for i in range(len(terms)):
        curve = Curve(term=terms[i], ytm=ytm[i] / 100)
        curve_list.append(curve)
    return curve_list


def plot_zero_curve(original_curve_list: List[Curve], zero_curve_list: List[Curve]):
    """绘制即期收益率曲线（黑白风格）"""
    terms = [curve.term for curve in zero_curve_list]
    zero_rates = [curve.zero_rate * 100 for curve in zero_curve_list]  # 转换为百分比
    original_terms = [original_curve.term for original_curve in original_curve_list]
    original_rates = [original_curve.zero_rate * 100 for original_curve in original_curve_list]
    plt.figure(figsize=(10, 6))
    plt.scatter(original_terms, original_rates, color="black", label="债券即期收益率")
    plt.plot(terms, zero_rates, linestyle="-", color="black", label="拟合后的收益率曲线", linewidth=2)
    plt.xlabel("期限（年）", fontsize=14)
    plt.ylabel("即期收益率（%）", fontsize=14)
    # 设置坐标轴颜色为黑色
    ax = plt.gca()
    ax.spines["bottom"].set_color("black")
    ax.spines["left"].set_color("black")
    ax.tick_params(axis="x", colors="black")
    ax.tick_params(axis="y", colors="black")
    # 去掉网格线
    ax.grid(False)
    plt.show()

def plot_ytm_curve(original_curve_list: List[Curve], zero_curve_list: List[Curve], circle_terms = [13,17,25]):
    """绘制到期收益率曲线（黑白风格）"""
    original_terms = [original_curve.term for original_curve in original_curve_list]
    original_rates = [original_curve.ytm * 100 for original_curve in original_curve_list]
    # circle_rates = [curve.ytm * 100 for curve in zero_curve_list if curve.term in circle_terms]
    plt.figure(figsize=(10, 6))
    plt.scatter(original_terms, original_rates, color="black", label="真实的债券到期收益率")
    # plt.scatter(circle_terms, circle_rates, color="black", label="缺失点", marker='o', facecolors="none", s=100)
    terms = [curve.term for curve in zero_curve_list]
    zero_rates = [curve.ytm * 100 for curve in zero_curve_list]  # 转换为百分比
    plt.plot(terms, zero_rates, linestyle="-", color="black", label="拟合后的收益率曲线", linewidth=2)
    terms = [curve.term for curve in zero_curve_list if curve.term not in original_terms]
    zero_rates = [curve.ytm * 100 for curve in zero_curve_list if curve.term not in original_terms]  # 转换为百分比
    plt.scatter(terms, zero_rates, color="black", label="补全后的到期收益率", marker='o', facecolors="none")
    plt.xlabel("期限（年）", fontsize=14)
    plt.ylabel("到期收益率（%）", fontsize=14)
    plt.legend()
    # 设置坐标轴颜色为黑色
    ax = plt.gca()
    ax.spines["bottom"].set_color("black")
    ax.spines["left"].set_color("black")
    ax.tick_params(axis="x", colors="black")
    ax.tick_params(axis="y", colors="black")
    # 去掉网格线
    ax.grid(False)
    plt.show()

#%%
from copy import deepcopy
# 付息频率
frequence = 2
# 根据收益率曲线编制零息利率曲线
curve_list = test_curve_list()
original_curve_list = deepcopy(curve_list)

#%% 三个点
curve_list_three = curve_list[4:7]
original_curve_three_list = deepcopy(curve_list_three)
curve_list_three_copy = completion_curve(curve_list_three, list_term(curve_list_three[0].term, curve_list_three[-1].term, 1, 3))
plot_ytm_curve(original_curve_three_list, curve_list_three_copy)
#%% 完整曲线
curve_list_copy = completion_curve(curve_list, list_term(1, 50, frequence, 3))
plot_ytm_curve(original_curve_list, curve_list_copy)

# %%
